# base_model.py
import logging
from typing import List, Optional, Iterable, Any

import pandas as pd
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.declarative import declarative_base

from infrastructure.database.database_context import DatabaseContext
from utils.logger import ConsoleLogger

Base = declarative_base()
logger = ConsoleLogger(name="BaseModel", level=logging.INFO)


class BaseModel(Base):
    """所有实体类的基类，提供通用的数据库操作方法"""

    __abstract__ = True

    @classmethod
    def save(cls, df: pd.DataFrame, if_exists: str = 'append',
             index_elements: Optional[List[str]] = None) -> bool:
        """保存DataFrame数据"""
        if df.empty:
            return True

        try:
            # 处理数据类型转换
            for col in df.select_dtypes(include=['datetime64[ns]']).columns:
                df[col] = pd.to_datetime(df[col], utc=True)

            records = df.to_dict('records')
            session = DatabaseContext.get_session()

            try:
                if index_elements and if_exists == 'append':
                    stmt = insert(cls.__table__).values(records)
                    stmt = stmt.on_conflict_do_nothing(
                        index_elements=index_elements
                    )
                    session.execute(stmt)
                else:
                    session.bulk_insert_mappings(cls, records)
                session.commit()
                return True
            except Exception:
                session.rollback()
                raise
        except Exception as e:
            logger.error(f"保存数据错误: {e}")
            return False

    @classmethod
    def find_all(cls) -> List['BaseModel']:
        """获取所有记录"""
        return DatabaseContext.get_session().query(cls).all()

    @classmethod
    def find_by_id(cls, id: Any) -> Optional['BaseModel']:
        """根据ID获取单条记录"""
        return DatabaseContext.get_session().query(cls).get(id)

    @classmethod
    def find_all_by_id(cls, ids: Iterable[Any]) -> List['BaseModel']:
        """根据ID列表获取记录"""
        return DatabaseContext.get_session().query(cls).filter(cls.id.in_(ids)).all()

    @classmethod
    def save_model(cls, instance: 'BaseModel', on_duplicate: str = 'ignore') -> 'BaseModel':
        """保存单个实例"""
        session = DatabaseContext.get_session()
        try:
            if on_duplicate in ['update', 'ignore']:
                existing = session.query(cls).get(instance.id)
                if existing:
                    if on_duplicate == 'update':
                        for key, value in instance.__dict__.items():
                            if not key.startswith('_'):
                                setattr(existing, key, value)
                        session.commit()
                        return existing
                    else:  # ignore
                        return existing

            if not sa.inspect(instance).persistent:
                instance = session.merge(instance)
            else:
                session.add(instance)
            session.commit()
            return instance
        except SQLAlchemyError as e:
            session.rollback()
            logger.error(f"Error saving {cls.__name__}: {str(e)}")
            raise

    @classmethod
    def update(cls, instance: 'BaseModel') -> None:
        """保存单个实例"""
        session = DatabaseContext.get_session()
        try:
            session.add(instance)
            session.commit()
        except SQLAlchemyError as e:
            session.rollback()
            logger.error(f"Error update {cls.__name__}: {str(e)}")
            raise

    @classmethod
    def delete(cls, instance: 'BaseModel') -> None:
        """删除单个实例"""
        session = DatabaseContext.get_session()
        try:
            if not sa.inspect(instance).persistent:
                instance = session.merge(instance)
            session.delete(instance)
            session.commit()
        except SQLAlchemyError as e:
            session.rollback()
            logger.error(f"Error deleting {cls.__name__}: {str(e)}")
            raise

    @classmethod
    def delete_by_id(cls, id: Any) -> None:
        """根据ID删除记录"""
        session = DatabaseContext.get_session()
        try:
            instance = session.query(cls).get(id)
            if instance:
                session.delete(instance)
                session.commit()
        except SQLAlchemyError as e:
            session.rollback()
            logger.error(f"Error deleting {cls.__name__} by id: {str(e)}")
            raise

    @classmethod
    def delete_all_by_id(cls, ids: Iterable[Any]) -> None:
        """根据ID列表批量删除"""
        session = DatabaseContext.get_session()
        try:
            instances = session.query(cls).filter(cls.id.in_(ids)).all()
            for instance in instances:
                session.delete(instance)
            session.commit()
        except SQLAlchemyError as e:
            session.rollback()
            logger.error(f"Error batch deleting {cls.__name__}: {str(e)}")
            raise

    @classmethod
    def to_dataframe(cls, results: List['BaseModel']) -> pd.DataFrame:
        """将查询结果转换为DataFrame"""
        records = []
        for obj in results:
            record = {}
            for col in obj.__table__.columns:
                value = getattr(obj, col.name)
                if isinstance(value, (sa.Numeric, sa.DECIMAL)):
                    value = float(value) if value is not None else None
                elif isinstance(value, sa.Integer):
                    value = int(value) if value is not None else None
                record[col.name] = value
            records.append(record)
        return pd.DataFrame.from_records(records)